// client.go - 客户端代码
package main

import (
	"fmt"
	"net"
	"os"
	"strconv"
	"time"
)

func main() {
	if len(os.Args) != 4 {
		fmt.Println("Usage: client <server-ip> <server-port> <local-port>")
		os.Exit(1)
	}

	serverIP := os.Args[1]
	serverPort := os.Args[2]
	localPort := os.Args[3]

	// 解析服务器地址
	serverAddr, err := net.ResolveUDPAddr("udp", 
		net.JoinHostPort(serverIP, serverPort))
	if err != nil {
		panic(err)
	}

	// 绑定本地端口
	localAddr, err := net.ResolveUDPAddr("udp", 
		net.JoinHostPort("0.0.0.0", localPort))
	if err != nil {
		panic(err)
	}

	conn, err := net.ListenUDP("udp", localAddr)
	if err != nil {
		panic(err)
	}
	defer conn.Close()

	// 向中继服务器注册
	_, err = conn.WriteToUDP([]byte("REGISTER"), serverAddr)
	if err != nil {
		panic(err)
	}

	fmt.Printf("Client registered on port %s\n", localPort)

	// 接收对端地址
	peerAddr := make(chan *net.UDPAddr, 1)
	go func() {
		// 数据缓冲区
		buffer := make([]byte, 1024)
		for {
			n, addr, err := conn.ReadFromUDP(buffer)
			if err != nil {
				fmt.Println("Read error:", err)
				continue
			}

			if addr.Port == serverAddr.Port {
				peerStr := string(buffer[:n])
				fmt.Printf("Received peer address: %s\n", peerStr)
				
				parsedAddr, err := net.ResolveUDPAddr("udp", peerStr)
				if err != nil {
					fmt.Println("Invalid peer address:", err)
					continue
				}
				// 将地址发送到通道，通知主逻辑
				peerAddr <- parsedAddr
				return
			}
		}
	}()

	// 启动通信循环
	go func() {
		peer := <-peerAddr
		fmt.Printf("Starting direct communication with %s\n", peer)

		// 发送心跳包
		go func() {
			tick := time.Tick(1 * time.Second)
			for {
				select {
				case <-tick:
					msg := []byte("PONG " + strconv.Itoa(int(time.Now().Unix())))
					_, err := conn.WriteToUDP(msg, peer)
					if err != nil {
						fmt.Println("Send error:", err)
					}
				}
			}
		}()

		// 接收消息
		buffer := make([]byte, 1024)
		for {
			n, addr, err := conn.ReadFromUDP(buffer)
			if err != nil {
				fmt.Println("Read error:", err)
				continue
			}

			if addr.String() == peer.String() {
				fmt.Printf("Received from %s: %s\n", addr, string(buffer[:n]))
			}
		}
	}()

	// 保持程序运行
	select {}
}